import os
import argparse
import pandas as pd
import numpy as np
from NARM import inference_narm
from time import time
import warnings

# 忽略警告
warnings.filterwarnings("ignore")

# 解析命令行参数
parser = argparse.ArgumentParser()
parser.add_argument('--offline', action='store_true', help='是否为离线模式')
parser.add_argument('--epochs', type=int, default=1, help='训练的轮数')
opt = parser.parse_args()
print(opt)

# 加载训练数据
def load_data(train_paths):
    # 如果提供了两个训练数据文件，则加载并合并
    if len(train_paths) == 2:
        train_data1 = pd.read_csv(train_paths[0])
        train_data2 = pd.read_csv(train_paths[1])
        train_data = pd.concat([train_data1, train_data2], axis=0)
    else:
        train_data = pd.read_csv(train_paths[0])

    return train_data

if __name__ == '__main__':
    # 记录程序开始时间
    start_time = time()

    # 根据命令行参数选择数据路径
    if opt.offline:
        data_directory = 'offline_data'
    else:
        data_directory = '../tcdata'

    # 训练数据路径
    train_data_paths = ['{}/train_click_log.csv'.format(data_directory), 
                        '{}/testA_click_log.csv'.format(data_directory)]
    train_click_data = load_data(train_data_paths)

    # 测试数据路径
    test_data_path = '{}/testB_click_log.csv'.format(data_directory)
    test_click_data = pd.read_csv(test_data_path)

    # NARM 推理
    inference_narm(train_click_data, test_click_data, opt.epochs)
    print('[-] NARM推理结束。')

    # 记录程序结束时间并计算推理时长
    end_time = time()
    print('[+] NARM推理所用时间: {} 分钟\n'.format((end_time - start_time) / 60))
